package com.example.toy.engine.dataset;

import com.example.toy.core.dataset.BigCollection;
import com.example.toy.core.dataset.BigPairCollection;
import com.example.toy.core.dataset.Pair;
import com.example.toy.core.engine.EngineContext;
import com.example.toy.core.function.SerializableBiFunction;
import com.example.toy.core.function.SerializableFunction;
import com.example.toy.engine.spark.SparkEngineContext;
import lombok.Getter;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function2;

/**
 * @author shenb
 * @date 2021-12-20 22:05
 */
@Getter
public class SparkBigPairCollection<K,V> implements BigPairCollection<K,V> {
  private SparkEngineContext jsc;
  private JavaPairRDD<K,V> rdd;

  public SparkBigPairCollection(SparkEngineContext jsc,
      JavaPairRDD<K, V> rdd) {
    this.jsc = jsc;
    this.rdd = rdd;
  }

  @Override
  public EngineContext getContext() {
    return jsc;
  }

  @Override
  public <R> BigCollection<R> map(SerializableFunction<Pair<K, V>, R> serializableFunction) {
    JavaRDD<R> result = rdd.map(tuple -> {
      return serializableFunction.apply(new Pair<>(tuple._1, tuple._2));
    });
    return new SparkBigCollection<>(jsc, result);
  }

  @Override
  public BigCollection<V> reduceByKey(SerializableBiFunction<V, V, V> biFunction) {

    JavaPairRDD<K, V> result = rdd.reduceByKey(new Function2<V, V, V>() {
      @Override
      public V call(V v1, V v2) throws Exception {
        return biFunction.apply(v1, v2);
      }
    });
    JavaRDD<V> values = result.values();
    return new SparkBigCollection<>(jsc, values);
  }
}
